import os
from agent import agent_factory
from ril import factory_builder
import argparse
import torch

os.environ['SDL_VIDEODRIVER'] = 'dummy'  # 禁画面
os.environ['SDL_AUDIODRIVER'] = 'dsp'  # 禁声音


def get_args():
    parser = argparse.ArgumentParser()
    # common
    parser.add_argument('--task', type=str, default='Breakout-v0')
    parser.add_argument('--agent', type=str, default='dqn', help='rl algorithm, e.g. sac_rule, dqn')
    parser.add_argument('--seed', type=int, default=0)
    parser.add_argument('--buffer-size', type=int, default=40000)
    parser.add_argument('--lr', type=float, default=0.0001)
    parser.add_argument('--gamma', type=float, default=0.99)
    parser.add_argument('--epoch', type=int, default=100)
    parser.add_argument('--step-per-epoch', type=int, default=100000)
    parser.add_argument('--step-per-collect', type=int, default=10)
    parser.add_argument('--update-per-step', type=float, default=0.1)
    parser.add_argument('--batch-size', type=int, default=32)
    parser.add_argument('--training-num', type=int, default=1)
    parser.add_argument('--test-num', type=int, default=1)
    parser.add_argument('--logdir', type=str, default='log')
    parser.add_argument('--render', type=float, default=0.)
    parser.add_argument('--stop-at-goal', action='store_true',
                        help='Stop when test rew >= env.threshold')
    parser.add_argument('--frames-stack', type=int, default=4)
    parser.add_argument('--resume-path', type=str, default=None)
    parser.add_argument('--watch', action='store_true',
                        help='watch the play of pre-trained policy only')
    parser.add_argument('--linear-sizes', type=int,
                        nargs='*', default=[])
    parser.add_argument(
        '--device', type=str,
        default='cuda' if torch.cuda.is_available() else 'cpu')

    # off policy
    parser.add_argument('--eps-test', type=float, default=0.005)
    parser.add_argument('--eps-train', type=float, default=1.)
    parser.add_argument('--eps-train-final', type=float, default=0.05)
    parser.add_argument('--target-update-freq', type=int, default=500)
    parser.add_argument('--n-step', type=int, default=3)
    parser.add_argument('--episode-per-collect', type=int, default=10)
    parser.add_argument('--repeat-per-collect', type=int, default=2)
    parser.add_argument('--rew-norm', action='store_true')

    # ac
    parser.add_argument('--actor-lr', type=float, default=1e-4)
    parser.add_argument('--critic-lr', type=float, default=1e-3)
    parser.add_argument('--alpha-lr', type=float, default=3e-4)
    parser.add_argument('--tau', type=float, default=0.005)
    parser.add_argument('--alpha', type=float, default=0.05)
    parser.add_argument('--auto-alpha', action="store_true", default=False)
    parser.add_argument('--noise_std', type=float, default=1.2)
    # a2c
    parser.add_argument('--vf-coef', type=float, default=0.5)
    parser.add_argument('--ent-coef', type=float, default=0.01)
    parser.add_argument('--eps-clip', type=float, default=0.2)
    parser.add_argument('--max-grad-norm', type=float, default=0.5)
    parser.add_argument('--gae-lambda', type=float, default=0.95)
    parser.add_argument('--bound-action-method', type=str, default="clip")
    parser.add_argument('--lr-decay', type=int, default=True)
    parser.add_argument('--dual-clip', type=float, default=None)
    parser.add_argument('--value-clip', type=int, default=1)

    # RIL
    parser.add_argument('--ril', action='store_true', help='Use ril')
    parser.add_argument('--ril-p0', type=float, default=1., help='Init probability to take ril suggestion, '
                                                                  'default 1.0')
    parser.add_argument('--ril-gamma', type=float, default=0.99, help='Rule probability decay rate γ, '
                                                                       'pt = p0 * γ^t, '
                                                                       'default 0.99')
    parser.add_argument('--ril-den', type=int, default=2000, help='pt = p0 * γ^(t/den), '
                                                                   'default 2000')
    parser.add_argument('--ril-p-test', type=float, default=0, help='Rule probability in test')
    return parser.parse_args()


if __name__ == '__main__':
    args = get_args()
    if not args.ril:
        args.ril_p0 = 0
        args.ril_p_test = 0
    factories = factory_builder(args.task, args.agent)
    agent_factory[args.agent](*factories, args)
